package com.tianji.aigc.service.impl;

import cn.hutool.core.util.StrUtil;
import com.alibaba.dashscope.app.Application;
import com.alibaba.dashscope.app.ApplicationParam;
import com.alibaba.dashscope.app.ApplicationResult;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.JsonUtils;
import com.tianji.aigc.config.DashScopeProperties;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.IChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.utils.TokenContext;
import com.tianji.common.utils.UserContext;
import io.reactivex.Flowable;
import lombok.RequiredArgsConstructor;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
@RequiredArgsConstructor
@ConditionalOnProperty(prefix = "tj.ai", name = "chat-type", havingValue = "APP")
public class AppAgentChatService implements IChatService {

    private final DashScopeProperties dashScopeProperties;
    private final IChatSessionService chatSessionService;

    // 存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();
    // 输出结束的标记
    private static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();
    StringBuilder outputBuilder = new StringBuilder();

    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        String conversationId = IChatService.getConversationId(sessionId);
        String token = TokenContext.getToken();
        Long userId = UserContext.getUser();
        Map<String, Object> toolsMap = new HashMap<>();
        for (var tool : dashScopeProperties.getAppAgent().getTools()){
            toolsMap.put(tool, Map.of("user_token", token));
        }
        Map<String, Object> bizParams = Map.of("user_defined_tokens", toolsMap);

        ApplicationParam param = ApplicationParam.builder()
                .apiKey(dashScopeProperties.getKey())
                .appId(dashScopeProperties.getAppAgent().getId())
                .bizParams(JsonUtils.toJsonObject(bizParams))
                .prompt(question)
                .incrementalOutput(true)
                .sessionId(conversationId)
                .build();
        Application application = new Application();
        try {
            Flowable<ApplicationResult> result = application.streamCall(param);
            // 将Flowable 转化为 Flux 进行处理输出
            return Flux.from(result)
                    .doFirst(() -> GENERATE_STATUS.put(conversationId, true))
                    .doOnError(throwable -> GENERATE_STATUS.remove(conversationId))
                    .doOnComplete(() -> GENERATE_STATUS.remove(conversationId))
                    .takeWhile(applicationResult -> GENERATE_STATUS.getOrDefault(conversationId, false))
                    .map(applicationResult -> {
                        String text = applicationResult.getOutput().getText();
                        outputBuilder.append(text);
                        return ChatEventVO.builder()
                                .eventType(ChatEventTypeEnum.DATA.getValue())
                                .eventData(text)
                                .build();
                    })
                    .doFinally(signalType -> {
                        var content = StrUtil.format("""
                            ------------
                            USER:{} \n
                            ASSISTANT:{}
                            ------------
                            """, question, outputBuilder.toString());
                        this.chatSessionService.update(sessionId, content, userId);
                    })
                    .concatWith(Flux.just(STOP_EVENT));
        } catch (NoApiKeyException e) {
            throw new RuntimeException(e);
        } catch (InputRequiredException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void stop(String sessionId) {
        // 移除标记
        GENERATE_STATUS.remove(sessionId);
    }
}
